use std::collections::HashMap;

mod binary_tree;

pub struct BinaryIndexedTree {
    capacity: usize,
    tree: Vec<i32>,
}

impl BinaryIndexedTree {
    pub fn init(size: usize) -> BinaryIndexedTree {
        Self { capacity: size+1, tree: vec![0; size+1] }
    }

    pub fn new(nums: &Vec<i32>) -> BinaryIndexedTree {
        let mut bit = Self { capacity: nums.len()+1, tree: vec![0; nums.len()+1] };
        for i in 1..bit.capacity {
            bit.tree[i] += nums[i-1];
            for j in (i - Self::lowbit(i as i32) as usize)..(i-1) {
                bit.tree[i] += nums[j as usize]; // add up children
            }
        }
        return bit;
    }

    pub fn add(&mut self, mut index: i32, val: i32) {
        while (index as usize) < self.capacity {
            self.tree[index as usize] += val;
            index += Self::lowbit(index); // update father
        }
    }

    // return sum of [1..index]  one-indexed
    pub fn query(&mut self, mut index: i32) -> i32 {
        let mut sum = 0;
        while index >= 1 {
            sum += self.tree[index as usize];
            index -= Self::lowbit(index);
        }
        return sum;
    }

    // inclusive query, one-indexed, return sum of values in [left, right]
    pub fn range(&mut self, left: i32, right: i32) -> i32 {
        return self.query(right) - self.query(left-1);
    }

    fn lowbit(x: i32) -> i32 {
        x & (-x)
    }
}

// leetcode 327
#[allow(non_snake_case)]
pub fn count_range_sum(nums: Vec<i32>, lower: i32, upper: i32) -> i32 {
    let n = nums.len();
    // presum[k] = sum of (0, k), exclusive
    let (mut presum, mut allNums): (Vec<i64>, Vec<i64>) = (vec![0; n+1], Vec::with_capacity(3 * n + 1));
    allNums.push(0); // presum[0] == 0
    for (i, v) in nums.iter().enumerate() {
        presum[i+1] = presum[i] + *v as i64;
        allNums.append(&mut vec![presum[i+1], presum[i+1]-lower as i64, presum[i+1]-upper as i64]);
    }

    // 离散化处理，让值散落在从1开始的正整数
    allNums.sort();
    println!("{:?}", allNums);
    let mut k = 1;
    let mut kth: HashMap<i64, i32> = HashMap::new();
    kth.insert(allNums[0], k);
    for i in 1..allNums.len() {
        if allNums[i] != allNums[i-1] {
            k += 1;
            kth.insert(allNums[i], k);
        }
    }

    let mut bit = BinaryIndexedTree::init(k as usize);
    // println!("k = {}", k);
    bit.add(kth[&0], 1);  // kth[presum[0]]
    let mut res = 0;
    for j in 1..presum.len() {
        let (left, right) = (kth[&(presum[j]-upper as i64)], kth[&(presum[j]-lower as i64)]);
        println!("query range[{}, {}]", left, right);
        res += bit.range(left, right);
        println!("presum[{}], add kth[{}]= {}",j, presum[j], kth[&presum[j]]);
        bit.add(kth[&presum[j]], 1);
    }
    res
}

#[test]
fn test_count_range_sum() {
    println!("{}", count_range_sum(vec![-2,5,-1], -2, 2));
    // println!("{}", count_range_sum(vec![2147483647,-2147483648,-1,0], -1, 0)); // 用i64存储
    // println!("{}", count_range_sum(vec![-2,0,1,-1],1,4)); // allNums要包含presum[0] todo: why?
}

#[test]
fn test_binary_indexed_tree() {
    let mut bit = BinaryIndexedTree::new(&vec![-2, 3, 5, 1]);
    println!("{}", bit.range(1, 3));
    // for j in 0..0 {
    //     println!("{}", j); // print nothing
    // }
}